import { codeToHtmlCached } from "@renderer/lib/shiki-highlighter";
import type { ThemedToken } from "shiki/core";

export interface ShikiPreProperties {
  class?: string;
  style?: string;
  tabindex?: number;
}

export interface HighlightChunkResult {
  lines: Array<{
    tokens?: ThemedToken[];
  }>;
  recall: number;
}

class ShikiStreamService {
  private highlighterCache = new Map<string, unknown>();
  private tokenizerMap = new Map<string, unknown>();

  private parseShikiHtml(html: string): ThemedToken[][] {
    const lineTokens: ThemedToken[][] = [];

    if (!html) return lineTokens;

    const tempDiv = document.createElement("div");
    tempDiv.innerHTML = html;

    const codeElement = tempDiv.querySelector("pre code");
    if (!codeElement) return lineTokens;

    const lineElements = codeElement.querySelectorAll(".line");

    if (lineElements.length > 0) {
      lineElements.forEach((lineElement, lineIndex) => {
        const tokens: ThemedToken[] = [];
        const spans = lineElement.querySelectorAll("span");
        let offset = 0;

        if (spans.length === 0) {
          const textContent = lineElement.textContent || "";
          tokens.push({
            content: textContent,
            offset: 0,
          });
        } else {
          spans.forEach((span) => {
            const content = span.textContent || "";
            const style = span.getAttribute("style") || "";

            const colorMatch = style.match(/color:\s*([^;]+)/);
            const bgColorMatch = style.match(/background-color:\s*([^;]+)/);

            tokens.push({
              content,
              offset,
              color: colorMatch?.[1]?.trim(),
              bgColor: bgColorMatch?.[1]?.trim(),
            });
            offset += content.length;
          });
        }

        lineTokens[lineIndex] = tokens;
      });
    } else {
      const fullText = codeElement.textContent || "";
      const lines = fullText.split("\n");

      lines.forEach((lineText, lineIndex) => {
        lineTokens[lineIndex] = [
          {
            content: lineText,
            offset: 0,
          },
        ];
      });
    }

    return lineTokens;
  }

  async highlightCodeChunk(
    trunk: string,
    language: string,
    theme: string,
    _callerId: string,
  ): Promise<HighlightChunkResult> {
    try {
      const themeMode =
        theme.includes("dark") || theme.includes("Dark") ? "dark" : "light";
      const html = await codeToHtmlCached(
        trunk,
        language,
        themeMode as "light" | "dark",
      );
      const lineTokens = this.parseShikiHtml(html);

      return {
        lines: lineTokens.map((tokens) => ({ tokens })),
        recall: 0,
      };
    } catch {
      const lines = trunk.split("\n");
      return {
        lines: lines.map(() => ({ tokens: [] })),
        recall: 0,
      };
    }
  }

  async highlightStreamingCode(
    fullContent: string,
    language: string,
    theme: string,
    _callerId: string,
  ): Promise<HighlightChunkResult> {
    try {
      const themeMode =
        theme.includes("dark") || theme.includes("Dark") ? "dark" : "light";
      const html = await codeToHtmlCached(
        fullContent,
        language,
        themeMode as "light" | "dark",
      );
      const lineTokens = this.parseShikiHtml(html);

      return {
        lines: lineTokens.map((tokens) => ({ tokens })),
        recall: 0,
      };
    } catch {
      const lines = fullContent.split("\n");
      return {
        lines: lines.map(() => ({ tokens: [] })),
        recall: 0,
      };
    }
  }

  async getShikiPreProperties(
    _language: string,
    theme: string,
  ): Promise<ShikiPreProperties> {
    const isDark = theme.includes("dark") || theme.includes("Dark");
    return {
      class: `shiki ${isDark ? "shiki-dark" : "shiki-light"}`,
      style: `background-color: hsl(var(--background)); color: hsl(var(--foreground));`,
      tabindex: 0,
    };
  }

  cleanupTokenizers(callerId: string): void {
    this.tokenizerMap.delete(callerId);
  }

  dispose(): void {
    this.highlighterCache.clear();
    this.tokenizerMap.clear();
  }
}

export const shikiStreamService = new ShikiStreamService();
